import inspect
from pickle import Unpickler, Pickler

class PickledExecution(object):
    """Wraps execution of a Python generator function with
    extra support code that pickles all yielded values
    into stdout, and pickles all raised exceptions to stderr"""
    
    def _get_final_code(self):
        # glue the code together with support code
        final_code = '''
import sys
from StringIO import StringIO

original_stdout = sys.stdout
original_stderr = sys.stderr
sys.stdout = StringIO()
sys.stderr = StringIO()

from pickle import Pickler, Unpickler

stdout_pickler = Pickler(original_stdout)
stderr_pickler = Pickler(original_stderr)
stdin_unpickler = Unpickler(sys.stdin)

%(generator_code)s

if __name__ == '__main__':
    
    try:
        if %(generator_name)s.func_code.co_argcount == 0:
            the_generator = %(generator_name)s()
        elif %(generator_name)s.func_code.co_argcount == 1:
            the_generator = %(generator_name)s(stdin_unpickler.load)
        else:
            raise TypeError("function does not have 0 or 1 arguments")

        for item in the_generator:
            stdout_pickler.dump(item)
        
    except Exception, e:
        exc_info = sys.exc_info()
        exc_type = exc_info[0]
        exc_value = exc_info[1]
        exception_metadata = (e, (exc_type, exc_value, None))
        stderr_pickler.dump( exception_metadata )
''' % dict(generator_name = self.__generator_name, generator_code = self.__generator_code)
    
        # run a syntax check on the final code
        compile(final_code, "<string>", "exec")
    
        # the code is clean, return it
        return final_code
    
    code = property(_get_final_code)
    """the actual code that can be executed"""
    
    def __init__(self, generator):
        """given a Python generator function, creates a chunk of Python
        code (as a string), which uses stdout to communicate pickled versions
        of each yielded value, and stderr to communicate a pickled version
        of any exception that gets raised.  This code can be accessed
        via the `code` property of the PickledExecution object"""
    
        # get the raw python code
        raw_code = inspect.getsource(generator)
        # docs = inspect.getdoc(generator) or ""
        # comments = inspect.getcomments(generator) or ""
        # just_code = raw_code.replace(docs, "").replace(comments, "")
    
        # remove bad indentation in the raw code
        raw_code_lines = raw_code.split("\n")
        raw_code_indent_prefix = ""
        for character in raw_code_lines[0]:
            if character.isspace():
                raw_code_indent_prefix += character
            else:
                break
        unindented_lines = []
        for line in raw_code_lines:
            unindented_line = line.replace(raw_code_indent_prefix, "", 1)
            unindented_lines.append( unindented_line )
        unindented_raw_code = "\n".join(unindented_lines)
    
        # run a syntax check on the code
        compile(unindented_raw_code, "<string>", "exec")
        
        # store the information we need about this code
        self.__generator_code = unindented_raw_code
        self.__generator_name = generator.__name__
    
    def write_object_to_stdin(self, obj, stdin):
        """given an object and a stdin file, write a pickled object to the stream"""
        stdin_pickler = Pickler(stdin)
        stdin_pickler.dump(obj)
    
    def get_response_iterator(self, stdout, stderr):
        """given stdout/stderr associated with pickled execution, iterates
        over all unpickled objects from stdout, and raise any exceptions
        from unpickling stderr"""
        
        # create some unpicklers to read the remote pickle output
        stdout_unpickler = Unpickler(stdout)
        stderr_unpickler = Unpickler(stderr)
    
        # for each item we got, yield it
        try:
            while True:
                item = stdout_unpickler.load()
                yield item
        except EOFError:
            pass
        except Exception, e:
            raise RuntimeError( "stdout: %s, stderr: %s" % (stdout.read(), stderr.read()) )
    
        # if we had an exception, yield that
        exc = None
        exc_info = None
        try:
            exc, exc_info = stderr_unpickler.load()
        except EOFError:
            exc = None
        except Exception, e:
            raise RuntimeError( "stdout: %s, stderr: %s" % (stdout.read(), stderr.read()) )
    
        # re-raise serialized exceptions
        if exc:
            raise create_wrapped_exception( exc, exc_info )
